import numpy as np
from audio_processing.base_processor import AudioProcessor
import threading
import time
from collections import deque


class FFTViewerProcessor(AudioProcessor):
    """FFT频谱查看器处理器"""

    def __init__(self, parameters):
        super().__init__(parameters)
        self.auto_scale = parameters.get('auto_scale', True)
        self.fft_size = int(parameters.get('fft_size', 1024))
        self.window_type = parameters.get('window_type', 'hann')
        self.magnitude_scale = parameters.get('magnitude_scale', 'log')
        self.frequency_range = parameters.get('frequency_range', 22050)
        self.frequency_scale = parameters.get('frequency_scale', 'linear')
        self.show_grid = parameters.get('show_grid', True)
        self.refresh_rate = parameters.get('refresh_rate', 5)  # 刷新率（FPS）

        # FFT数据缓存 - 根据FFT大小设置缓冲区
        self.max_buffer_size = self.fft_size * 2  # 保留两帧数据以便平滑
        self.audio_buffer = deque(maxlen=self.max_buffer_size)
        self.sample_count = 0

        # 显示窗口引用
        self.display_window = None
        self.node_name = "频谱查看器"

        # 线程锁
        self.data_lock = threading.Lock()

        # 处理状态
        self.is_processing = False
        self.last_update_time = 0
        self.update_interval = 1.0 / self.refresh_rate  # 根据刷新率计算更新间隔

        # 参数更新回调
        self.parameter_update_callback = None
    
    def _get_window_function(self):
        """获取窗函数"""
        size = self.fft_size
        if self.window_type == 'hann':
            return np.hanning(size)
        elif self.window_type == 'hamming':
            return np.hamming(size)
        elif self.window_type == 'blackman':
            return np.blackman(size)
        else:  # rectangular
            return np.ones(size)

    def _calculate_fft(self, audio_data):
        """计算FFT并返回频率和幅度数据"""
        # 如果数据不足FFT大小，填充零
        if len(audio_data) < self.fft_size:
            padded_data = np.zeros(self.fft_size)
            padded_data[:len(audio_data)] = audio_data
            audio_data = padded_data
        else:
            # 取最近的FFT大小的数据
            audio_data = audio_data[-self.fft_size:]
        
        # 应用窗函数
        window = self._get_window_function()
        windowed_data = audio_data * window
        
        # 计算FFT（只取实数FFT的一半）
        fft_result = np.fft.rfft(windowed_data)
        
        # 计算频率轴
        frequencies = np.fft.rfftfreq(self.fft_size, 1.0 / self.sample_rate)
        
        # 计算幅度
        magnitude = np.abs(fft_result)
        
        # 根据选择的幅度刻度进行转换
        if self.magnitude_scale == 'log':
            # 转换为dB，添加小值避免log(0)
            magnitude = 20 * np.log10(magnitude + 1e-10)
        
        # 根据频率范围过滤
        freq_mask = frequencies <= self.frequency_range
        frequencies = frequencies[freq_mask]
        magnitude = magnitude[freq_mask]
        
        return frequencies, magnitude
    
    def set_sample_rate(self, sample_rate: int):
        """设置采样率"""
        super().set_sample_rate(sample_rate)

    def process(self, audio_data: np.ndarray) -> np.ndarray:
        """处理音频数据并缓存用于FFT计算"""
        if not self.is_active:
            return audio_data

        # 确保数据形状正确
        audio_data = self._ensure_correct_shape(audio_data)

        # 缓存音频数据用于FFT
        with self.data_lock:
            # 将音频数据添加到缓冲区，处理多声道（取平均值）
            for i in range(len(audio_data)):
                if audio_data.shape[1] > 0:
                    # 对于多声道，取平均值
                    sample = np.mean(audio_data[i])
                else:
                    sample = 0.0
                self.audio_buffer.append(sample)
                self.sample_count += 1

        # 限制更新频率
        current_time = time.time()
        if current_time - self.last_update_time < self.update_interval:
            return audio_data

        self.last_update_time = current_time

        # 更新显示窗口（如果存在）
        if self.display_window and hasattr(self.display_window, 'update_fft_signal'):
            try:
                # 获取当前音频数据并计算FFT
                with self.data_lock:
                    buffer_data = list(self.audio_buffer)
                    
                    # 只有当缓冲区有足够数据时才计算FFT
                    if len(buffer_data) >= self.fft_size:
                        frequencies, magnitude = self._calculate_fft(buffer_data)
                    else:
                        frequencies, magnitude = [], []
                    
                    current_params = {
                        'auto_scale': self.auto_scale,
                        'fft_size': self.fft_size,
                        'window_type': self.window_type,
                        'magnitude_scale': self.magnitude_scale,
                    'frequency_range': self.frequency_range,
                    'frequency_scale': self.frequency_scale,
                    'show_grid': self.show_grid,
                        'sample_rate': self.sample_rate
                    }

                # 在UI线程中更新频谱
                self.display_window.update_fft_signal.emit((frequencies.tolist() if len(frequencies) > 0 else []), 
                                                          (magnitude.tolist() if len(magnitude) > 0 else []), 
                                                          current_params)

            except Exception as e:
                print(f"更新频谱显示时出错: {e}")
                import traceback
                traceback.print_exc()

        # FFT查看器通常是直通的（不影响音频信号）
        return audio_data

    def update_parameters(self, new_parameters):
        """更新参数"""
        super().update_parameters(new_parameters)

        # 更新参数，添加适当的类型转换和范围检查
        self.auto_scale = new_parameters.get('auto_scale', self.auto_scale)
        
        # FFT大小
        new_fft_size = new_parameters.get('fft_size', self.fft_size)
        try:
            # 确保是有效的FFT大小选项
            valid_sizes = [256, 512, 1024, 2048, 4096]
            fft_size = int(new_fft_size)
            if fft_size in valid_sizes:
                self.fft_size = fft_size
                # 更新缓冲区大小
                self.max_buffer_size = self.fft_size * 2
                with self.data_lock:
                    old_data = list(self.audio_buffer)
                    self.audio_buffer = deque(maxlen=self.max_buffer_size)
                    self.audio_buffer.extend(old_data)
        except (ValueError, TypeError):
            pass
        
        # 窗函数类型
        self.window_type = new_parameters.get('window_type', self.window_type)
        
        # 幅度刻度
        self.magnitude_scale = new_parameters.get('magnitude_scale', self.magnitude_scale)
        
        # 频率范围
        new_freq_range = new_parameters.get('frequency_range', self.frequency_range)
        
        # 频率刻度
        self.frequency_scale = new_parameters.get('frequency_scale', self.frequency_scale)
        try:
            freq_range = float(new_freq_range)
            # 限制在合理范围内
            self.frequency_range = max(1000, min(freq_range, self.sample_rate // 2))
        except (ValueError, TypeError):
            pass
        
        self.show_grid = new_parameters.get('show_grid', self.show_grid)
        
        # 更新刷新率参数
        if 'refresh_rate' in new_parameters:
            try:
                refresh_rate = int(new_parameters['refresh_rate'])
                # 限制在有效范围内
                self.refresh_rate = max(5, min(refresh_rate, 20))
                # 更新更新间隔
                self.update_interval = 1.0 / self.refresh_rate
            except (ValueError, TypeError):
                pass

        # 通知显示窗口参数已更新
        if self.display_window and hasattr(self.display_window, 'on_processor_parameters_updated'):
            self.display_window.on_processor_parameters_updated(new_parameters)

    def set_display_window(self, window):
        """设置显示窗口"""
        self.display_window = window

    def set_node_name(self, name):
        """设置节点名称"""
        self.node_name = name

    def get_fft_data(self):
        """获取当前FFT数据"""
        with self.data_lock:
            buffer_data = list(self.audio_buffer)
            if len(buffer_data) >= self.fft_size:
                frequencies, magnitude = self._calculate_fft(buffer_data)
                return frequencies.tolist(), magnitude.tolist(), self.sample_count
        return [], [], self.sample_count

    def reset(self):
        super().reset()
        with self.data_lock:
            self.audio_buffer.clear()
            self.sample_count = 0

    def start(self):
        """启动处理器"""
        super().start()
        self.is_processing = True

    def stop(self):
        """停止处理器"""
        super().stop()
        self.is_processing = False

    def set_parameter_update_callback(self, callback):
        """设置参数更新回调"""
        self.parameter_update_callback = callback

    def update_parameter_from_display(self, param_name, value):
        """从显示窗口接收参数更新"""
        self.update_parameters({param_name: value})